import pickle
import gradio as gr
import numpy as np
from PIL import Image
from scipy.ndimage import median_filter, gaussian_filter

# 加载KNN模型
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

# 简化但高效的数字识别预测函数
def predict_digit(sketch_image):
    try:
        # 基本检查
        if sketch_image is None:
            return "请绘制数字"
        
        # 处理Gradio Sketchpad返回的字典格式数据
        if isinstance(sketch_image, dict):
            if 'layers' in sketch_image and len(sketch_image['layers']) > 0:
                sketch_image = sketch_image['layers'][0]
            elif 'background' in sketch_image:
                sketch_image = sketch_image['background']
            else:
                return "无法识别的图像格式"
        
        # 确保是numpy数组
        if not isinstance(sketch_image, np.ndarray):
            sketch_image = np.array(sketch_image)
        
        # 处理RGB/RGBA图像
        if len(sketch_image.shape) == 3:
            if sketch_image.shape[2] == 4:  # RGBA格式
                # 使用Alpha通道创建掩码
                alpha_channel = sketch_image[:, :, 3]
                white_background = np.ones((sketch_image.shape[0], sketch_image.shape[1]), dtype=np.uint8) * 255
                black_foreground = np.zeros((sketch_image.shape[0], sketch_image.shape[1]), dtype=np.uint8)
                sketch_image = np.where(alpha_channel > 128, black_foreground, white_background)
            elif sketch_image.shape[2] >= 3:  # RGB格式
                # 转换为灰度图像
                sketch_image = 0.299 * sketch_image[:, :, 0] + 0.587 * sketch_image[:, :, 1] + 0.114 * sketch_image[:, :, 2]
        
        # 确保图像是2D
        if len(sketch_image.shape) != 2:
            return "图像格式错误"
        
        # 确保是uint8类型
        if sketch_image.dtype != np.uint8:
            sketch_image = sketch_image.astype(np.uint8)
        
        # 颜色反转 - 确保白字黑底
        white_pixels = np.sum(sketch_image > 200)
        total_pixels = sketch_image.size
        white_ratio = white_pixels / total_pixels
        
        if white_ratio > 0.5:
            sketch_image = 255 - sketch_image
        
        # 关键改进：使用双阶段大小调整以保留更多细节
        img = Image.fromarray(sketch_image)
        # 第一阶段：调整到20x20（比直接到8x8保留更多细节）
        img = img.resize((20, 20), Image.Resampling.LANCZOS)
        # 第二阶段：从20x20调整到8x8
        img = img.resize((8, 8), Image.Resampling.LANCZOS)
        img_array = np.array(img)
        
        # 高级噪声过滤
        # 1. 中值滤波去除孤立点
        img_array = median_filter(img_array, size=2)
        # 2. 高斯模糊平滑边缘
        img_array = gaussian_filter(img_array, sigma=0.6)
        
        # 增强对比度 - 使用自适应拉伸
        non_zero_values = img_array[img_array > 5]
        if len(non_zero_values) > 0:
            # 使用10%和90%百分位数以避免异常值
            min_val = np.percentile(non_zero_values, 10)
            max_val = np.percentile(non_zero_values, 90)
            
            if max_val > min_val:
                # 对比度拉伸
                img_array = 255 * ((img_array - min_val) / (max_val - min_val))
                img_array = np.clip(img_array, 0, 255).astype(np.uint8)
        
        # 智能阈值处理 - 针对不同数字特点调整
        mean_val = np.mean(img_array)
        std_val = np.std(img_array)
        
        # 根据图像特性动态调整阈值
        if std_val < 30:  # 低对比度图像
            threshold = int(max(10, mean_val - 15))
        else:  # 高对比度图像
            threshold = int(max(15, mean_val - std_val / 4))
        
        img_array = np.where(img_array > threshold, img_array, 0).astype(np.uint8)
        
        # 数字居中处理 - 关键改进
        non_zero = np.nonzero(img_array)
        if len(non_zero[0]) > 0:
            # 找到数字边界
            min_row, max_row = np.min(non_zero[0]), np.max(non_zero[0])
            min_col, max_col = np.min(non_zero[1]), np.max(non_zero[1])
            
            # 裁剪到数字区域
            digit_only = img_array[min_row:max_row+1, min_col:max_col+1]
            
            # 计算中心偏移
            h, w = digit_only.shape
            pad_h = (8 - h) // 2
            pad_w = (8 - w) // 2
            
            # 创建居中图像
            centered = np.zeros((8, 8), dtype=np.uint8)
            centered[pad_h:pad_h+h, pad_w:pad_w+w] = digit_only
            img_array = centered
        
        # 标准化到MNIST格式 (0-16)
        if np.max(img_array) > 0:
            img_array = (img_array / 255.0 * 16.0).astype(np.uint8)
        else:
            img_array = np.zeros_like(img_array, dtype=np.uint8)
        
        # 检查是否为空图像
        if np.sum(img_array) < 1:
            return "请绘制数字"
        
        # 准备预测输入
        img_flat = img_array.flatten().reshape(1, -1)
        
        # 模型预测
        prediction = knn_model.predict(img_flat)
        
        # 返回简洁的预测结果
        return str(int(prediction[0]))
    except Exception as e:
        return "识别错误"

# 清除函数
def clear_sketch():
    return None, ""

# 构建极简界面
with gr.Blocks(title="数字识别") as demo:
    gr.Markdown("## 手写数字识别")
    
    with gr.Row():
        with gr.Column():
            sketch = gr.Sketchpad(label="绘图区域")
            with gr.Row():
                clear = gr.Button("清除")
                submit = gr.Button("识别", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="识别结果", lines=1, placeholder="识别结果将显示在这里")
    
    # 绑定事件
    submit.click(
        fn=predict_digit,
        inputs=sketch,
        outputs=output
    )
    
    clear.click(
        fn=clear_sketch,
        inputs=None,
        outputs=[sketch, output]
    )
    
    # 添加简单的CSS以优化界面
    demo.css = """
    .gr-sketchpad { height: 300px; }
    .gr-button-primary { background-color: #2563eb; }
    .gr-textbox { font-size: 24px; text-align: center; }
    """

# 启动服务
if __name__ == "__main__":
    demo.launch(share=True, server_name="127.0.0.1")